package com.hzw.saas.service.hpc.service.impl;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.extra.ssh.JschUtil;
import cn.hutool.json.JSONUtil;
import com.hzw.saas.api.hpc.bo.HpcAccountProviderBo;
import com.hzw.saas.api.property.IPropertyService;
import com.hzw.saas.api.property.enmus.SysDictKeyEnum;
import com.hzw.saas.common.config.util.AssertUtil;
import com.hzw.saas.service.hpc.service.IHpcAccountSlurmService;
import com.hzw.saas.service.hpc.service.bo.SlurmAccountBO;
import com.hzw.saas.service.hpc.util.HpcUtils;
import com.jcraft.jsch.Session;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.nio.charset.Charset;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;
import java.util.stream.Collectors;

/**
 * @author sonam
 * @sine 2021/9/13 2:10 下午
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class HpcAccountSlurmServiceImpl implements IHpcAccountSlurmService {
    @Resource(name = IPropertyService.SYS_DICT_BEAN)
    private final IPropertyService<SysDictKeyEnum> propertyService;

    @Override
    public List<SlurmAccountBO> getSlurmAccounts() {
        return getSlurmAccountDTOS();
    }

    @Override
    public SlurmAccountBO getSlurmAccount(String slurmAccountName) {
        try {
            // 查询Slurm用户
            List<SlurmAccountBO> slurmAccounts = getSlurmAccountDTOS();
            slurmAccounts = slurmAccounts.stream().filter(slurmAccount -> Objects.equals(slurmAccount.getSlurmAccount(), slurmAccountName)).collect(Collectors.toList());
            SlurmAccountBO slurmAccount = null;
            if (CollectionUtil.isNotEmpty(slurmAccounts)) {
                slurmAccount = slurmAccounts.get(0);
            }
            return slurmAccount;
        } catch (Exception e) {
            return null;
        }
    }

    private List<SlurmAccountBO> getSlurmAccountDTOS() {
        String propertyValue = propertyService.getPropertyValue(SysDictKeyEnum.SLURM_ACCOUNTS);
        return JSONUtil.toList(propertyValue, SlurmAccountBO.class);
    }

    @Override
    public void addSlurmUser(HpcAccountProviderBo providerBo, String user, String slurmAccount) {
        Session session = getSession(providerBo);
        try {
            this.addSlurmUser(session, user, slurmAccount);
        } finally {
            JschUtil.close(session);
        }
    }

    private void addSlurmUser(Session session, String user, String slurmAccount) {
        try {
            SlurmAccountBO slurmAccountBo = getSlurmAccount(slurmAccount);
            AssertUtil.assertThrow("添加slurmUser失败，请重试", Objects.isNull(slurmAccountBo));
            AssertUtil.assertThrow("参数错误, slurm账号不能为空", StrUtil.isBlank(slurmAccount));
            // 先删除，不删除会导致无法更新
            this.deleteSlurmUser(session, user, slurmAccount);
            execCommands(session, StrUtil.format("sudo sacctmgr -i add user name={} account={}", user, slurmAccount));
            execCommands(session, StrUtil.format("sudo sacctmgr -i modify  user {} set  Grptres=\"node={}\" MaxJobs={} MaxSubmitJobs={}", user, slurmAccountBo.getNode(), slurmAccountBo.getMaxJobs(), slurmAccountBo.getMaxSubmitJobs()));
            AssertUtil.assertThrow("添加slurmUser失败，请重试", !exist(user, slurmAccount, session));
        } finally {
            JschUtil.close(session);
        }
    }

    @Override
    public void deleteSlurmUser(HpcAccountProviderBo providerBo, String user, String slurmAccount) {
        Session session = getSession(providerBo);
        try {
            this.deleteSlurmUser(session, user, slurmAccount);
        } finally {
            JschUtil.close(session);
        }
    }

    private void deleteSlurmUser(Session session, String user, String slurmAccount) {
        // 删除slurm用户，将所有pending任务全部退出
        try {
            List<String> jobIds = HpcUtils.pullJobIdByStatus(session, user, "PENDING");
            String command = jobIds.stream().reduce("", (a, b) -> a.concat("sudo scancel ").concat(b).concat(StrUtil.SPACE).concat(StrUtil.LF));
            execCommands(session, command);
        } catch (Exception e) {
            // 忽略退出任务，带来的异常
            log.error("退出等待任务失败", e);
        }
        execCommands(session, StrUtil.format("sudo sacctmgr -i delete user name={}", user));
        AssertUtil.assertThrow("删除slurmUser失败，请重试", exist(user, slurmAccount, session));
    }

    @Override
    public void updateSlurmUser(HpcAccountProviderBo providerBo, String user, String slurmAccount) {
        Session session = getSession(providerBo);
        try {
            this.deleteSlurmUser(session, user, slurmAccount);
            this.addSlurmUser(session, user, slurmAccount);
        } finally {
            JschUtil.close(session);
        }
    }

    private Session getSession(HpcAccountProviderBo providerBo) {
        return JschUtil.openSession(providerBo.getAdminHost(), providerBo.getAdminPort(), providerBo.getAdminUserName(), providerBo.getAdminPassword());
    }

    private String execCommands(Session session, String command) {
        log.info("\n{}", command);
        return JschUtil.exec(session, command, Charset.defaultCharset());
    }

    private boolean exist(String user, String slurmAccount, Session session) {
        String slurmUserStr = execCommands(session, "sudo sacctmgr list user");
        String[] slurmUsers = slurmUserStr.split(StrUtil.LF);
        for (int i = 2; i < slurmUsers.length; i++) {
            String slurmUserLine = slurmUsers[i];
            List<String> lines = Arrays.stream(slurmUserLine.split(StrUtil.SPACE)).filter(StrUtil::isNotBlank).map(StrUtil::trim).collect(Collectors.toList());
            if (user.length() > 10) {
                user = user.substring(0, 9).concat("+");
            }
            if (slurmAccount.length() > 10) {
                slurmAccount = slurmAccount.substring(0, 9).concat("+");
            }
            if (Objects.equals(lines.get(0), user) && (Objects.equals(lines.get(1), slurmAccount) || StrUtil.isBlank(slurmAccount))) {
                return true;
            }
        }
        return false;
    }

//    public static void main(String[] args) {
//        IHpcAccountSlurmService slurmService = new HpcAccountSlurmServiceImpl(null);
//        HpcAccountProviderBo providerBo = new HpcAccountProviderBo();
//        providerBo.setAdminHost("60.216.39.134");
//        providerBo.setAdminPort(22022);
//        providerBo.setAdminUserName("admin");
//        providerBo.setAdminPassword("Zhu@88jie");
//        slurmService.addSlurmUser(providerBo, "aaaaaa", "aaaaa");
//    }
}
